import json
import re
from itertools import chain
from typing import Annotated, Any, List

import ngram
import uvicorn
from cachetools import TTLCache
from fastapi import FastAPI, Form
from pydantic import BaseModel, field_validator

app = FastAPI()

with open("./assets/level_1.txt") as f:
    all_level_1 = f.read().strip().split("\n")
    all_level_1_ngram = ngram.NGram(all_level_1)

with open("./assets/level_2.txt") as f:
    all_level_2 = f.read().strip().split("\n")
    all_level_2_ngram = ngram.NGram(all_level_2)

with open("./assets/tag2category.json") as f:
    scene2category = json.load(f)
    all_scene = list(scene2category.keys())

all_scene_ngram = ngram.NGram(scene2category.keys())
scene_cache = TTLCache(maxsize=128, ttl=10 * 60 * 60)


def remove_common_elements(*sets):
    if sets:
        # 找出所有集合的交集
        common = set.intersection(*sets)

        # 从每个集合中移除共有元素
        return set.union(*[s - common for s in sets])
    else:
        return set()


class CachedModel(BaseModel):
    conversation_id: str = ""
    level_1: str = ""
    level_2: str = ""
    scene: List[str] = []
    legacy_scene: List[str] = []
    keywords: List[str] = []
    irrelevant: List[str] = []
    # 已经引导的场景类别，避免重复引导
    categories: List[str] = []
    search_results: Any = []

    @field_validator("level_1")
    def validate_level_1(cls, value):
        return value.strip()

    @field_validator("level_2")
    def validate_level_2(cls, value):
        return value.strip()

    @field_validator("scene", mode="before")
    def validate_scene(cls, value):
        if not isinstance(value, list):
            value = [value]
        results = []

        for item in value:
            item = item.strip().replace('"', "")
            if not value:
                continue
            item = re.split(r"[-，]", item)
            results.extend(item)
        # 硬编码，省内跨省纠正
        for i, item in enumerate(results):
            if "跨省" in item and "省内" in item:
                results[i] = results[i].replace("跨省", "省内")

        return [item for item in results if item]

    @field_validator("keywords", mode="before")
    def validate_keywords(cls, value):
        if not isinstance(value, list):
            value = [value]

        results = []
        for item in value:
            item = item.strip().replace('"', "")
            if not item:
                continue
            item = re.split(r"[-，]", item)
            results.extend(item)
        return [item for item in results if item]

    @field_validator("categories", mode="before")
    def validate_categories(cls, value):
        if isinstance(value, list):
            return value

        value = value.strip()
        if not value:
            return []
        if value.count("，") > value.count("-"):
            return value.split("，")
        else:
            return value.split("-")

    @field_validator("search_results", mode="before")
    def validate_search_results(cls, value):
        if isinstance(value, str):
            value = json.loads(value)

        for item in value:
            content = item["content"]
            lines = content.split("\n")
            item["_tags"] = lines[0].split("-")
            item["_question"] = lines[1]

        return value

    @field_validator("irrelevant", mode="before")
    def validate_irrelevant(cls, value):
        if not isinstance(value, list):
            value = [value]

        results = []
        for item in value:
            item = item.strip().replace('"', "")
            if not item:
                continue
            item = re.split(r"[-，]", item)
            results.extend(item)
        return [item for item in results if item]

    def merge(self, other):
        self.level_1 = other.level_1 if other.level_1 else self.level_1
        self.level_2 = other.level_2 if other.level_2 else self.level_2
        self.legacy_scene = self.scene
        self.scene = [item for item in set(self.scene) | set(other.scene) if item]

        self.keywords = other.keywords
        self.irrelevant = other.irrelevant
        self.categories = [
            item for item in set(self.categories) | set(other.categories) if item
        ]
        self.search_results = other.search_results

    def to_dict(self):
        # hardcoding here 这里社保卡不继承老的场景
        if self.level_1 == "社保卡":
            # scene = self.legacy_scene
            scene = []
        else:
            scene = self.scene

        if self.level_1 and self.level_2:
            elements = [self.level_1, self.level_2] + scene
        else:
            elements = [self.level_1, self.level_2] + scene + self.keywords
        elements = [item for item in elements if item]
        return {
            "level_1": self.level_1,
            "level_2": self.level_2,
            "scene": "-".join(self.scene),
            "keywords": "-".join(self.scene),
            "search_query": "-".join(elements),
            "irrelevant": "-".join(self.irrelevant),
        }


@app.post("/graph")
async def predict(cached_model: Annotated[CachedModel, Form()]) -> dict:
    # 0 引导一级事项
    # -1 超纲拒答

    print(cached_model)
    if not cached_model.level_1 and (cached_model.irrelevant or cached_model.keywords):
        return {"candidates": [], "flag": -1}
    # 硬编码
    elif cached_model.irrelevant and "保" in "".join(cached_model.irrelevant):
        return {"candidates": all_level_1, "flag": -1}
    elif cached_model.keywords and "保" in "".join(cached_model.keywords):
        return {"candidates": all_level_1, "flag": -1}
    # 硬编码

    elif cached_model.level_1 == "社保":
        return {"candidates": all_level_1, "flag": 0}
    elif not cached_model.level_1 and (cached_model.level_2 or cached_model.scene):
        return {"candidates": all_level_1, "flag": 0}
    else:
        return {"candidates": [], "flag": -1}


@app.post("/slots")
def slots(cached_model: Annotated[CachedModel, Form()]):
    conversation_id = cached_model.conversation_id
    if conversation_id not in scene_cache:
        scene_cache[conversation_id] = cached_model
    else:
        scene_cache[conversation_id].merge(cached_model)

    return scene_cache[conversation_id].to_dict()


@app.post("/clear")
def slots(conversation_id: str):
    if conversation_id not in scene_cache:
        return {"message": "No avaliable session found."}
    cached_model = scene_cache[conversation_id]
    # cached_model.level_1 = ""
    # cached_model.level_2 = ""
    return {"message": "Success"}


@app.post("/query")
def query(cached_model: Annotated[CachedModel, Form()]):

    print(scene_cache[cached_model.conversation_id])
    level_1 = cached_model.level_1
    level_2 = cached_model.level_2
    if not level_2 and not cached_model.keywords:
        return {
            "num": 0,
            "filtered_search_results": [],
            "need_reask": "True",
            "candidates": f"询问用户具体办理{cached_model.level_1}的什么业务",
            "scores": [],
            "reason": "",
        }
    conversation_id = cached_model.conversation_id
    scene = cached_model.scene
    search_results = cached_model.search_results
    scores = [item["metadata"]["score"] for item in search_results]
    max_score = max(scores) if scores else 0

    confident_results = [
        item
        for item in search_results
        if max_score - item["metadata"]["score"]
        < max(1 - max_score, 0.03)  # 过滤掉第二梯队分值的
        and item["_tags"][0] in level_1  # 过滤掉一级事项不符合的
    ]

    keywords = list(chain(item["_tags"] for item in confident_results))

    candidates = remove_common_elements(*[set(k) for k in keywords])
    candidates = [item for item in candidates if item in all_scene]

    categories = {}
    for candidate in candidates:
        c = all_scene_ngram.find(candidate)
        if c:
            category = scene2category[c]
            if category not in categories:
                categories[category] = []
            categories[category].append(candidate)

    for item in scene:
        c = all_scene_ngram.find(item)
        if c:
            categories.pop(scene2category[c], None)

    for item in scene_cache[conversation_id].categories:
        categories.pop(item, None)

    category = max(categories, key=lambda x: len(categories[x])) if categories else ""
    candidates = categories[category] if category else []
    if category:
        scene_cache[conversation_id].categories.append(category)

    if len(candidates) <= 1:
        candidates = []

    return {
        "num": len(confident_results),
        "filtered_search_results": confident_results,
        "need_reask": "False" if not candidates else "True",
        "candidates": "，".join(candidates) if len(candidates) > 1 else "",
        "scores": scores,
        "reason": "",
    }


if __name__ == "__main__":
    uvicorn.run(app, port=8889)
